from langchain_openai import ChatOpenAI
import spacy
from langchain.schema import BaseMemory
#from langchain_core.pydantic_v1 import BaseModel
from langchain.chains.llm import LLMChain
from typing import Any,Iterable,Dict,List,Optional
from langchain.prompts import PromptTemplate
from pydantic import BaseModel

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-3.5-turbo",api_key=api_key ,base_url=api_base)

nlp = spacy.load("zh_core_web_sm")
text = "2014年9月，马云在杭州发起了世界互联网大会。这一事件标志着中国在全球互联网发展中的积极参与。"
doc = nlp(text)
#print(doc.ents)

text = "随后，2015年11月，腾讯公司在深圳发布了微信支付功能，进一步推动了中国数字经济的发展"
doc = nlp(text)
#print(doc.ents)

text = "与此同时，字节跳动的产品——抖音，在2016年成为中国年轻人中的热门应用，为内容创造者提供了展示才华的平台。"
doc = nlp(text)
#print(doc.ents)

class SpacyEntityMemory(BaseMemory,BaseModel):
    """存储实体信息的记忆类。"""

    # 定义一个字典来存储实体的信息。
    entities:dict={}

    # 定义一个 Key，作用是将有关实体的信息传递到提示符：memory_key
    memory_key:str="entities"

    # 清除全部的实体信息
    def clear(self):
        self.entities={}

    @property
    def memory_variables(self)->List[str]:
        """定义提供给提示符的变量。"""
        return [self.memory_key]

    def save_context(self, inputs: dict[str, Any], outputs: dict[str, str]) -> None:
        """将此对话中的上下文保存到缓冲区。"""
        # 获取输入文本
        print("inputs:",inputs)
        inputKey=inputs.keys()
        print("inputKey:",inputKey)
        listKey=list(inputKey)[0]
        print("listKey:",listKey)
        text = inputs[listKey]
        print("text:",text)

        # 运行spacy，从输入文本中提取关键的信息
        doc = nlp(text)
        print(doc)
        print("doc.ents:",doc.ents)
        for ent in doc.ents:
            ent_str = str(ent)
            # 如果实体信息已经存在，将当前文本（text）追加到这个实体键下已有的文本后面，每个文本之间用换行符\n分隔。

            # 示例：{'2014年9月': '2014年9月，马云在杭州发起了世界互联网大会。这一事件标志着中国在全球互联网发展中的积极参与。', \
            #       '马云': '2014年9月，马云在杭州发起了世界互联网大会。这一事件标志着中国在全球互联网发展中的积极参与。',
            #      '杭州': '2014年9月，马云在杭州发起了世界互联网大会。这一事件标志着中国在全球互联网发展中的积极参与。',
            #      '中国': '2014年9月，马云在杭州发起了世界互联网大会。这一事件标志着中国在全球互联网发展中的积极参与。'}
            if ent_str in self.entities:
                self.entities[ent_str] += f"\n{text}"
            else:
                self.entities[ent_str] = text




    def load_memory_variables(self,inputs:Dict[str,Any])->Dict[str,Any]:
        # 获取输入文本
        inputKey=inputs.keys()
        print(inputKey)
        listKey = list(inputKey)
        print(listKey)
        doc = nlp(inputs[listKey[0]])
        print(doc)
        # 提取关于实体的已知信息(如果存在的话)。
        entities = [
            self.entities[str(ent)] for ent in doc.ents if str(ent) in self.entities
        ]

        # 返回要放在上下文中的实体的组合信息。
        return {self.memory_key:"\n".join(entities)}

template = """以下是人类和AI之间的友好对话。AI很健谈，并提供了大量来自其上下文的具体细节。如果AI不知道某个问题的答案，它会真诚地说不知道。如果相关，你会得到人类提到的实体的信息。

相关实体信息：
{entities}

对话：
人类: {input}
AI:"""

prompt = PromptTemplate(template=template,input_variables=["entities", "input"])
spacyMemory = SpacyEntityMemory(memory_key="entities")
chain = LLMChain(llm=llm,
                 prompt=prompt,
                 memory=spacyMemory,
                 verbose=True)
chain.invoke({"input":"2014年9月，马云在杭州发起了世界互联网大会。这一事件标志着中国在全球互联网发展中的积极参与。"})

chain.invoke({"input":"随后，2015年11月，腾讯公司在深圳发布了微信支付功能，进一步推动了中国数字经济的发展"})
chain.invoke({"input":"那天，我正好去游玩杭州，碰到了马云"})
response=chain.invoke({"input":"杭州西湖旁边的雷峰塔"})
print(response)